import os
import numpy as np
import torch
from typing import List, Tuple, Dict, Set, Optional
import heapq
import random # Added for random.choice
import time

from agent.dqn_agent import DQNAgent
from env.grid_env import GridMap


class DQNPlanner:
    """
    使用训练好的DQN模型进行单机器人路径规划
    作为CBS算法的低层规划器
    """
    def __init__(self, grid_map: GridMap, model_path_or_agent, device: str = None):
        """
        初始化DQN规划器
        
        Args:
            grid_map: 网格地图
            model_path_or_agent: DQN模型路径或已加载的DQNAgent实例
            device: 计算设备
        """
        self.grid_map = grid_map
        
        # 检查是否传入的是已加载的智能体
        if isinstance(model_path_or_agent, DQNAgent):
            self.agent = model_path_or_agent
            print("DQN规划器使用已加载的智能体")
        else:
            # 传入的是模型路径
            model_path = model_path_or_agent
            self.model_path = model_path
            
            # 创建DQN智能体
            state_dim = 15  # 更新后的状态维度: 10(基本) + 1(振荡特征) + 4(上一动作的one-hot编码)
            action_dim = 4  # [上, 右, 下, 左]
            
            # 简单配置，只用于推理
            config = {
                'learning_rate': 0.001,
                'gamma': 0.99
            }
            
            self.agent = DQNAgent(
                state_dim=state_dim,
                action_dim=action_dim,
                config=config,
                device=device
            )
            
            # 加载模型
            if os.path.exists(model_path):
                self.agent.load_model(model_path)
                print(f"DQN规划器已加载模型: {model_path}")
            else:
                raise FileNotFoundError(f"模型文件 {model_path} 不存在!")
        
        # 设置规划超时时间（秒）
        self.planning_timeout = 5.0
    
    def _is_stuck_in_corner(self, position, path, threshold=3):
        """
        检测机器人是否卡在角落或边界
        
        Args:
            position: 当前位置 (x, y)
            path: 已经走过的路径
            threshold: 判断为卡住的重复次数阈值
            
        Returns:
            bool: 是否卡在角落
        """
        # 如果路径太短，不判断
        if len(path) < threshold * 2:
            return False
            
        # 检查最近的几个位置是否重复出现
        recent_positions = path[-threshold*2:]
        position_counts = {}
        
        for pos in recent_positions:
            if pos in position_counts:
                position_counts[pos] += 1
            else:
                position_counts[pos] = 1
                
        # 如果有位置重复出现超过阈值，判断为卡住
        for pos, count in position_counts.items():
            if count >= threshold:
                # 检查是否在边界或角落
                x, y = pos
                if (x == 0 or x == self.grid_map.width-1 or 
                    y == 0 or y == self.grid_map.height-1):
                    return True
                    
                # 检查是否被障碍物包围（至少两个方向有障碍）
                obstacles = 0
                if not self.grid_map.is_valid_position((x-1, y)): obstacles += 1
                if not self.grid_map.is_valid_position((x+1, y)): obstacles += 1
                if not self.grid_map.is_valid_position((x, y-1)): obstacles += 1
                if not self.grid_map.is_valid_position((x, y+1)): obstacles += 1
                
                if obstacles >= 2:
                    return True
                    
        return False
        
    def _get_escape_action(self, current_pos, goal, path):
        """
        当机器人卡住时，获取逃脱动作
        
        Args:
            current_pos: 当前位置 (x, y)
            goal: 目标位置 (x, y)
            path: 已经走过的路径
            
        Returns:
            int: 逃脱动作
        """
        x, y = current_pos
        
        # 获取所有有效动作
        valid_actions = []
        for action in range(len(self.grid_map.ACTIONS)):
            dx, dy = self.grid_map.ACTIONS[action]
            next_pos = (x + dx, y + dy)
            if self.grid_map.is_valid_position(next_pos):
                valid_actions.append(action)
                
        if not valid_actions:
            return 0  # 如果没有有效动作，返回默认动作
            
        # 找出之前没走过的动作
        recent_positions = set(path[-5:])  # 最近5个位置
        new_actions = []
        
        for action in valid_actions:
            dx, dy = self.grid_map.ACTIONS[action]
            next_pos = (x + dx, y + dy)
            if next_pos not in recent_positions:
                new_actions.append(action)
                
        if new_actions:
            # 随机选择一个新动作
            return random.choice(new_actions)
        else:
            # 如果所有动作都走过了，随机选择一个有效动作
            return random.choice(valid_actions)

    def plan_path(self, start: Tuple[int, int], goal: Tuple[int, int], 
                  constraints: Set[Tuple[int, int, int]] = None, 
                  max_timesteps: int = 100) -> List[Tuple[int, int]]:
        """
        使用DQN规划从起点到目标的路径
        
        Args:
            start: 起点 (x, y)
            goal: 目标 (x, y)
            constraints: 约束集合 (x, y, time_step)，表示在特定时间步不能到达的位置
            max_timesteps: 最大时间步
            
        Returns:
            List[Tuple[int, int]]: 规划的路径
        """
        path = [start]
        current_pos = start
        stuck_count = 0
        consecutive_same_pos = 0
        last_pos = None
        
        # 将约束转换为字典格式，便于查询
        constraints_dict = {}
        if constraints:
            for x, y, t in constraints:
                if t not in constraints_dict:
                    constraints_dict[t] = set()
                constraints_dict[t].add((x, y))
        
        # 使用DQN逐步规划路径
        for timestep in range(max_timesteps):
            if current_pos == goal:
                break
            
            # 检查是否卡住
            is_stuck = self._is_stuck_in_corner(current_pos, path)
            
            # 检查是否连续多次在同一位置
            if current_pos == last_pos:
                consecutive_same_pos += 1
            else:
                consecutive_same_pos = 0
                
            last_pos = current_pos
            
            # 如果卡住或连续多次在同一位置，尝试逃脱
            if is_stuck or consecutive_same_pos >= 3:
                action = self._get_escape_action(current_pos, goal, path)
                stuck_count += 1
                
                # 如果连续多次卡住，尝试更激进的逃脱策略
                if stuck_count > 5:
                    # 尝试找到一个远离当前位置的有效位置
                    valid_positions = []
                    for y in range(self.grid_map.height):
                        for x in range(self.grid_map.width):
                            if self.grid_map.is_valid_position((x, y)) and (x, y) != current_pos:
                                # 计算与当前位置的距离
                                dist = abs(x - current_pos[0]) + abs(y - current_pos[1])
                                if dist > 3:  # 至少要离开3步
                                    valid_positions.append((x, y))
                    
                    if valid_positions:
                        # 找到一个中间目标点
                        mid_goal = min(valid_positions, key=lambda pos: abs(pos[0] - goal[0]) + abs(pos[1] - goal[1]))
                        
                        # 获取状态向量
                        state = self.grid_map.get_state_vector(current_pos, mid_goal)
                        
                        # 选择动作（评估模式）
                        action = self.agent.select_action(state, eval_mode=True)
            else:
                # 获取状态向量
                state = self.grid_map.get_state_vector(current_pos, goal)
                
                # 选择动作（评估模式）
                action = self.agent.select_action(state, eval_mode=True)
            
            # 计算下一个位置
            dx, dy = self.grid_map.ACTIONS[action]
            next_pos = (current_pos[0] + dx, current_pos[1] + dy)
            
            # 检查位置是否有效
            if not self.grid_map.is_valid_position(next_pos):
                # 如果无效，尝试其他动作
                valid_actions = []
                for a in range(self.agent.action_dim):
                    dx, dy = self.grid_map.ACTIONS[a]
                    pos = (current_pos[0] + dx, current_pos[1] + dy)
                    if self.grid_map.is_valid_position(pos):
                        valid_actions.append(a)
                
                if valid_actions:
                    # 从有效动作中选择Q值最大的
                    with torch.no_grad():
                        state_tensor = torch.FloatTensor(state).unsqueeze(0).to(self.agent.device)
                        q_values = self.agent.policy_net(state_tensor).cpu().numpy()[0]
                        # 只考虑有效动作的Q值
                        valid_q_values = [(q_values[a], a) for a in valid_actions]
                        _, action = max(valid_q_values)
                        
                        dx, dy = self.grid_map.ACTIONS[action]
                        next_pos = (current_pos[0] + dx, current_pos[1] + dy)
                else:
                    # 如果没有有效动作，原地等待
                    next_pos = current_pos
            
            # 检查约束条件
            if timestep + 1 in constraints_dict and next_pos in constraints_dict[timestep + 1]:
                # 如果违反约束，尝试其他动作
                valid_actions = []
                for a in range(self.agent.action_dim):
                    dx, dy = self.grid_map.ACTIONS[a]
                    pos = (current_pos[0] + dx, current_pos[1] + dy)
                    if self.grid_map.is_valid_position(pos) and (
                        timestep + 1 not in constraints_dict or 
                        pos not in constraints_dict[timestep + 1]
                    ):
                        valid_actions.append(a)
                
                if valid_actions:
                    # 从有效动作中选择Q值最大的
                    with torch.no_grad():
                        state_tensor = torch.FloatTensor(state).unsqueeze(0).to(self.agent.device)
                        q_values = self.agent.policy_net(state_tensor).cpu().numpy()[0]
                        # 只考虑有效动作的Q值
                        valid_q_values = [(q_values[a], a) for a in valid_actions]
                        _, action = max(valid_q_values)
                        
                        dx, dy = self.grid_map.ACTIONS[action]
                        next_pos = (current_pos[0] + dx, current_pos[1] + dy)
                else:
                    # 如果没有有效动作，原地等待
                    next_pos = current_pos
            
            # 更新当前位置
            current_pos = next_pos
            path.append(current_pos)
            
            # 如果到达目标，结束规划
            if current_pos == goal:
                break
        
        return path
    
    def plan_path_astar(self, start: Tuple[int, int], goal: Tuple[int, int], 
                         constraints: Set[Tuple[int, int, int]] = None, 
                         max_timesteps: int = 100) -> List[Tuple[int, int]]:
        """
        使用A*算法规划路径（作为备选方案，当DQN无法找到有效路径时使用）
        
        Args:
            start: 起点 (x, y)
            goal: 终点 (x, y)
            constraints: 约束集合 (x, y, timestep)
            max_timesteps: 最大时间步数
            
        Returns:
            List[Tuple[int, int]]: 路径点列表
        """
        if constraints is None:
            constraints = set()
        
        # 将约束转换为字典
        constraints_dict = {}
        for x, y, t in constraints:
            if t not in constraints_dict:
                constraints_dict[t] = set()
            constraints_dict[t].add((x, y))
        
        # 启发式函数：曼哈顿距离
        def heuristic(pos):
            return abs(pos[0] - goal[0]) + abs(pos[1] - goal[1])
        
        # 初始化开放列表和关闭列表
        open_list = []
        closed_set = set()
        
        # 使用优先队列实现开放列表
        # 格式: (f_score, timestep, position, path)
        heapq.heappush(open_list, (heuristic(start), 0, start, [start]))
        
        while open_list and len(open_list) < 10000:  # 防止无限循环
            # 获取f值最小的节点
            _, timestep, current, path = heapq.heappop(open_list)
            
            # 如果到达目标
            if current == goal:
                return path
            
            # 如果超过最大时间步数
            if timestep >= max_timesteps:
                continue
            
            # 如果节点已经访问过
            if (current, timestep) in closed_set:
                continue
            
            # 标记为已访问
            closed_set.add((current, timestep))
            
            # 遍历所有可能的动作
            for action in range(len(self.grid_map.ACTIONS)):
                dx, dy = self.grid_map.ACTIONS[action]
                next_pos = (current[0] + dx, current[1] + dy)
                next_timestep = timestep + 1
                
                # 检查位置是否有效
                if not self.grid_map.is_valid_position(next_pos):
                    continue
                
                # 检查是否违反约束
                if next_timestep in constraints_dict and next_pos in constraints_dict[next_timestep]:
                    continue
                
                # 如果节点未访问过
                if (next_pos, next_timestep) not in closed_set:
                    # 计算f值
                    g_score = next_timestep  # g值就是时间步
                    h_score = heuristic(next_pos)
                    f_score = g_score + h_score
                    
                    # 添加到开放列表
                    new_path = path + [next_pos]
                    heapq.heappush(open_list, (f_score, next_timestep, next_pos, new_path))
        
        # 如果找不到路径，返回空列表
        return []
    
    def plan_multiple_paths(self, start: Tuple[int, int], goal: Tuple[int, int], 
                       constraints: Set[Tuple[int, int, int]] = None, 
                       num_paths: int = 5, max_timesteps: int = 100) -> List[Tuple[int, int]]:
        """
        尝试生成多条路径，并返回最好的一条
        
        Args:
            start: 起点 (x, y)
            goal: 目标 (x, y)
            constraints: 约束集合 (x, y, time_step)
            num_paths: 生成的路径数量
            max_timesteps: 最大时间步
            
        Returns:
            List[Tuple[int, int]]: 最佳路径
        """
        start_time = time.time()
        print(f"开始为 {start} 到 {goal} 规划多条路径")
        
        # 先尝试A*算法
        astar_path = self.plan_path_astar(start, goal, constraints, max_timesteps)
        if astar_path:
            print(f"A*找到路径，长度: {len(astar_path)}")
        
        paths = []
        
        # 确保至少有一条路径
        if astar_path:
            paths.append(astar_path)
        
        # 尝试DQN规划多条路径
        for i in range(num_paths):
            # 检查是否超时
            if time.time() - start_time > self.planning_timeout:
                print(f"路径规划超时，已生成 {len(paths)} 条路径")
                break
                
            try:
                path = self.plan_path(start, goal, constraints, max_timesteps)
                if path and len(path) > 1:  # 确保路径有效且不只是起点
                    # 检查路径是否到达目标
                    if path[-1] == goal:
                        paths.append(path)
                        print(f"DQN生成路径 {i+1}/{num_paths}，长度: {len(path)}")
            except Exception as e:
                print(f"DQN规划路径时出错: {str(e)}")
        
        if not paths:
            # 如果没有找到任何路径，尝试简单的启发式搜索
            print("尝试回退到简单启发式搜索")
            fallback_path = self._fallback_search(start, goal, max_timesteps)
            if fallback_path:
                paths.append(fallback_path)
                print(f"启发式搜索找到路径，长度: {len(fallback_path)}")
        
        # 如果没有找到任何路径，返回空列表
        if not paths:
            print(f"无法找到从 {start} 到 {goal} 的路径")
            return []
        
        # 返回最短的路径
        shortest_path = min(paths, key=len)
        print(f"选择最短路径，长度: {len(shortest_path)}")
        return shortest_path
    
    def _fallback_search(self, start: Tuple[int, int], goal: Tuple[int, int], max_steps: int = 100) -> List[Tuple[int, int]]:
        """
        简单的回退搜索算法，当其他方法失败时使用
        
        Args:
            start: 起点 (x, y)
            goal: 目标 (x, y)
            max_steps: 最大步数
            
        Returns:
            List[Tuple[int, int]]: 规划的路径
        """
        # 使用广度优先搜索(BFS)作为最后的回退方案
        queue = [(start, [start])]  # (当前位置, 路径)
        visited = {start}
        
        while queue and len(visited) < max_steps:
            (x, y), path = queue.pop(0)
            
            if (x, y) == goal:
                return path
            
            # 尝试四个方向
            for dx, dy in self.grid_map.ACTIONS:
                nx, ny = x + dx, y + dy
                next_pos = (nx, ny)
                
                if self.grid_map.is_valid_position(next_pos) and next_pos not in visited:
                    visited.add(next_pos)
                    new_path = path + [next_pos]
                    queue.append((next_pos, new_path))
        
        return []

    def plan_path_with_fallback(self, start: Tuple[int, int], goal: Tuple[int, int], 
                                constraints: Set[Tuple[int, int, int]] = None, 
                                max_timesteps: int = 100) -> List[Tuple[int, int]]:
        """
        使用DQN规划路径，如果失败则回退到A*算法
        
        Args:
            start: 起点 (x, y)
            goal: 目标 (x, y)
            constraints: 约束集合 (x, y, time_step)
            max_timesteps: 最大时间步
            
        Returns:
            List[Tuple[int, int]]: 规划的路径
        """
        # 首先尝试使用DQN规划多条路径并选择最短的
        path = self.plan_multiple_paths(start, goal, constraints, num_paths=3, max_timesteps=max_timesteps)
        
        # 检查路径是否有效（到达目标）
        if path and path[-1] == goal:
            return path
        
        # 如果DQN失败，使用A*算法
        print(f"DQN规划失败，使用A*算法从 {start} 到 {goal}")
        path = self.plan_path_astar(start, goal, constraints, max_timesteps)
        
        return path if path else [start]  # 如果A*也失败，返回只包含起点的路径 